import logging
import os

from control_models.base import ControlModel, MODEL_ROOT, get_bool
from typing import List, Dict
import cv2
import numpy as np
from shapely.geometry import Polygon
from shapely.ops import unary_union
from utils.neural_nets import (
    BaseNN,
    MultiLabelLSTM,
    cached_feature_extraction,
    cached_yolo_predict,
)
from utils.converter import (
    get_label_map,
    convert_timelinelabels_to_probs,
    convert_probs_to_timelinelabels,
)
logger = logging.getLogger(__name__)


class PolygonLabelsModel(ControlModel):
    """
    Class representing a PolygonLabels control tag for YOLO model.
    """

    type = "PolygonLabels"
    model_path = "yolov8n-seg.pt"

    @classmethod
    def is_control_matched(cls, control,mode=None) -> bool:
        # check object tag type
        if control.objects[0].tag != "Image":
            return False
        return control.tag == cls.type and mode is None

    def predict_regions(self, path) -> List[Dict]:
        results = self.model.predict(path)
        return self.create_polygons(results, path)

    def create_polygons(self, results, path, target_sides=4, use_min_quadrilateral=False):
        logger.debug(f"create_polygons: {self.from_name}")
        data = results[0].masks  # take masks from the first frame
        model_names = self.model.names
        regions = []
        self.extra_params
        for i in range(len(data)):
            score = float(results[0].boxes.conf[i])  # tensor => float
            original_points = (
                data.xyn[i] * 100
            )  # get the polygon points for the current instance
            model_label = model_names[int(results[0].boxes.cls[i])]

            logger.debug(
                "----------------------\n"
                f"task id > {path}\n"
                f"type: {self.control}\n"
                f"polygon points > {original_points}\n"
                f"model label > {model_label}\n"
                f"score > {score}\n"
            )

            # bbox score is too low
            if score < self.model_score_threshold:
                continue

            # there is no mapping between model label and LS label
            if model_label not in self.label_map:
                continue
            output_label = self.label_map[model_label]

            # 如果需要拟合为指定边数的多边形
            if target_sides and target_sides > 2 and not use_min_quadrilateral:
                points = self.fit_polygon_to_sides(original_points, target_sides)
            
            # 如果需要拟合为最小四边形
            if use_min_quadrilateral:
                points = self.fit_min_bounding_quadrilateral(original_points)

            # 确保多边形点数量有效
            if len(points) < 3:
                logger.warning(f"多边形点数量不足，跳过: {len(points)}点")
                continue
            
            # add new region with polygon
            region = {
                "from_name": self.from_name,
                "to_name": self.to_name,
                "type": "polygonlabels",
                "value": {
                    "polygonlabels": [output_label],
                    "points": points.tolist(),  # Converting the tensor to a list for JSON serialization
                    "closed": True,
                },
                "score": score,
            }
            regions.append(region)
        return regions



    def fit(self, event, data, **kwargs):
        if not self.trainable:
            logger.debug(
                'TimelineLabels model is in not trainable mode. '
                'Use model_trainable="true" to enable training.'
            )
            return

        """Fit the model."""
        if event == "START_TRAINING":
            # TODO: the full training makes a lot of sense here, but it's not implemented yet
            raise NotImplementedError(
                f"The event START_TRAINING is not supported for this control model: {self.control.tag}"
            )

        if event in ("ANNOTATION_CREATED", "ANNOTATION_UPDATED"):
            features, labels, label_map, project_id = self.load_features_and_labels(
                data
            )
            classifier, path = self.load_polygons(features, label_map, project_id)
            return self.train_polygons(classifier, features, labels, path)
        
        
    def train_polygons(self, classifier, features, labels, path):
        """Train the classifier model for timelinelabels using incremental partial learning."""
        # Stop training when accuracy or f1 score reaches this threshold, it helps to avoid overfitting
        # because we partially train it on a small dataset from one annotation only
        get = self.control.attr.get
        epochs = int(
            get("model_polygons_epochs", 1000)
        )  # Maximum number of training epochs
        f1_threshold = float(get("model_polygons_f1_threshold", 0.95))
        accuracy_threshold = float(get("model_polygons_accuracy_threshold", 1.00))

        # Train and save
        result = classifier.partial_fit(
            features,
            labels,
            epochs=epochs,
            f1_threshold=f1_threshold,
            accuracy_threshold=accuracy_threshold,
        )
        classifier.save_and_cache(path)
        return result

    def load_polygons(self, features, label_map, project_id):
        """Load or create a classifier model for timelinelabels.
        1. Load neural network parameters from labeling config.
        2. Try loading classifier model from memory cache, then from disk.
        3. Or create a new classifier instance if there wasn't successful loading, or if parameters have changed.
        """
        get = self.control.attr.get
        # LSTM sequence size
        sequence_size = int(get("model_polygons_sequence_size", 16))
        # LSTM hidden state size
        hidden_size = int(get("model_polygons_hidden_size", 32))
        # LSTM num layers
        num_layers = int(get("model_polygons_num_layers", 1))

        # Load classifier
        path = self.get_polygons_path(project_id)
        classifier = BaseNN.load_cached_model(path)

        # Create a new classifier instance if it doesn't exist
        # or if labeling config has changed
        if (
            not classifier
            or classifier.label_map != label_map
            or classifier.sequence_size != sequence_size
            or classifier.hidden_size != hidden_size
            or classifier.num_layers != num_layers
        ):
            logger.info("Creating a new classifier model for timelinelabels")
            input_size = len(features[0])
            output_size = len(label_map)
            classifier = MultiLabelLSTM(
                input_size,
                output_size,
                sequence_size=sequence_size,
                hidden_size=hidden_size,
                num_layers=num_layers,
            )
            classifier.set_label_map(label_map)

        return classifier, path

    def load_features_and_labels(self, data):
        """Load features and labels from the annotation
        Args:
            data: event data, dictionary with keys 'task' and 'annotation'
        Returns:
            features: List of features, 2D array with shape (num_frames, num_features)
            labels: List of labels, 2D array with shape (num_frames, num_labels)
            label_map: Label map, dictionary mapping label names to indices in the labels array
            project_id: Project ID from Label Studio
        """
        # Get the task and regions from the annotation
        task = data["task"]
        project_id = task["project"]
        annotation = data["annotation"]
        regions = annotation["result"]

        # Get the features and labels for training
        video_path = self.get_path(task)
        frames = cached_feature_extraction(
            self.model, video_path, self.model.model_name
        )
        features = [frame.probs for frame in frames]
        label_map = get_label_map(self.control.labels)
        labels, used_labels = convert_timelinelabels_to_probs(
            regions, label_map=label_map, max_frame=len(frames)
        )

        # Check if all labels from used_labels are in the label_map
        if not used_labels.issubset(label_map.keys()):
            raise ValueError(
                f"Annotation labels set ({used_labels}) is not subset "
                f"of labels from the labeling config:\n{self.control}\n"
                f"It can be caused by the mismatch between the labeling config "
                f"and labels in the annotation #{data['annotation']['id']}"
                f"of project #{project_id}."
            )
        return features, labels, label_map, project_id

    def get_polygons_path(self, project_id):
        yolo_base_name = os.path.splitext(os.path.basename(self.model.model_name))[0]
        path = f"{MODEL_ROOT}/timelinelabels-{project_id}-{yolo_base_name}-{self.from_name}.pkl"
        return path
    
    def fit_polygon_to_sides(self, points, target_sides):
        """将多边形拟合为指定边数"""
        if len(points) <= target_sides:
            return points  # 如果点数已经小于等于目标边数，直接返回
        
        # 计算多边形周长，用于确定epsilon参数
        perimeter = cv2.arcLength(points, True)
        # 动态调整epsilon值以达到目标边数
        epsilon = perimeter * 0.01
        target_epsilon = epsilon
        best_approx = None
        best_diff = float('inf')
        
        # 二分法寻找最佳epsilon值
        for _ in range(20):  # 迭代次数
            approx = cv2.approxPolyDP(points, target_epsilon, True)
            approx_sides = len(approx)
            diff = abs(approx_sides - target_sides)
            
            if diff < best_diff:
                best_diff = diff
                best_approx = approx
                if diff == 0:
                    break  # 找到完美匹配
            
            # 调整epsilon值
            if approx_sides > target_sides:
                target_epsilon += epsilon * 0.1
            else:
                target_epsilon -= epsilon * 0.1
        
        return best_approx.reshape(-1, 2)

    def fit_min_bounding_quadrilateral(self, points):
        """计算最小外接四边形"""
        # 转换为Shapely多边形对象
        polygon = Polygon(points)
        
        # 计算最小外接矩形(可能是旋转的)
        min_rect = polygon.minimum_rotated_rectangle
        
        # 如果最小外接矩形是四边形，直接使用
        if len(min_rect.exterior.coords) == 5:  # 四边形有5个点(包括闭合点)
            return np.array(min_rect.exterior.coords[:-1])  # 排除最后一个闭合点
        
        # 否则使用凸包近似
        convex_hull = polygon.convex_hull
        # 对凸包进行四边形拟合
        perimeter = cv2.arcLength(np.array(convex_hull.exterior.coords), True)
        epsilon = perimeter * 0.02
        approx = cv2.approxPolyDP(np.array(convex_hull.exterior.coords), epsilon, True)
        
        # 确保结果是四边形
        if len(approx) != 4:
            # 如果不是四边形，使用最小面积矩形
            rect = cv2.minAreaRect(points.astype(np.float32))
            box = cv2.boxPoints(rect)
            return np.array(box)
        
        return approx.reshape(-1, 2)


# pre-load and cache default model at startup
PolygonLabelsModel.get_cached_model(PolygonLabelsModel.model_path)
